package cg.cmm.rag.config;

import cg.cmm.rag.service.Customer;
import cg.cmm.rag.tool.ToolUtils;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

/**
 * @author: baibing
 * create time: 2025-03-20 10:43
 * description:
 */
@Configuration
@RequiredArgsConstructor
public class CustomerInit {

    @Bean
    public Customer init(EmbeddingStore<TextSegment> embeddingStore){

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .timeout(Duration.ofSeconds(360))
                .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                .modelName("qwen2.5:latest") // 模型名称
                .apiKey("your-api-key") // 如果需要 API 密钥
                .build();

        // 手动创建EmbeddingStoreContentRetriever并设置embeddingModel
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .build();
        return AiServices.builder(Customer.class)
                .chatMemoryProvider(id -> MessageWindowChatMemory.withMaxMessages(1000))
                .contentRetriever(contentRetriever)
                .chatLanguageModel(OpenAiChatModel.builder()
                        .timeout(Duration.ofSeconds(360))
                        .baseUrl("http://localhost:11434/v1") // Ollama 本地服务地址
                        .modelName("qwen2.5:latest") // 模型名称
                        .apiKey("your-api-key") // 如果需要 API 密钥
                        .build())
                .tools(new ToolUtils())
                .build();
    }
}
